import os
import pandas as pd
from tqdm import tqdm

from config.config import config
import logger
from utils.read import read_pkl
from utils.write import write_pkl

log = logger.get_logger()


def main_select_cluster():
    dataDirPath = config["cluster"]["attrs_select_result_dir_path"]
    gtDataFileName = os.path.split(config["input"]["gt_data_file_path"])[1].split(".")[0]
    data = read_pkl(r"%s/%s.pkl" % (dataDirPath, gtDataFileName))
    lowerCount = config["cluster"]["lower_count"]
    upperCount = config["cluster"]["upper_count"]

    print("select personIds start....")
    # 1、去除样本数不在[lowerCount~upperCount]范围内的所有簇
    # 1.1、获取该列数据
    dataAttr = data["gt_person_id"]
    log.debug("this dataset have " + str(dataAttr.shape[0]) + " pieces of data")
    # 1.2、去重复
    log.info("remove duplicates...")
    dataTargetDedup = dataAttr.drop_duplicates()
    log.debug("this dataset have " + str(dataTargetDedup.shape[0]) + " clusters")
    log.info("remove duplicates end and select start...")
    # 1.3、建立一个map
    clustersMap = {cluster: 0 for cluster in dataTargetDedup}
    for target in list(dataAttr):
        clustersMap[target] += 1
    log.info("statistic end and sort start...")
    # 1.4 根据字典筛选簇
    personIdsMap = {}
    for key in clustersMap.keys():
        count = clustersMap.get(key)
        if lowerCount <= count <= upperCount:
            personIdsMap[key] = count

    log.info("cluster count: %d" % len(personIdsMap))
    # 2、在dataRes中选cluster_count个簇
    clusterCount = config["cluster"]["cluster_count"]
    selectMode = config["cluster"]["select_mode"]
    log.info("cluster %d cluster" % clusterCount)
    log.info("cluster mode : %s" % selectMode)

    if len(personIdsMap) < clusterCount:
        log.error("clusterCount is too large, curClusterCount is %d" % len(personIdsMap))
    selectPersonIds = []
    if "define".__eq__(selectMode):
        selectPersonIds = list(personIdsMap.keys())[:clusterCount]
    elif "max".__eq__(selectMode):
        # 优先选大簇
        for item in sorted(personIdsMap.items(), key=lambda item: item[1], reverse=True)[:clusterCount]:
            selectPersonIds.append(item[0])
    elif "min".__eq__(selectMode):
        # 优先选小簇
        for item in sorted(personIdsMap.items(), key=lambda item: item[1], reverse=False)[:clusterCount]:
            selectPersonIds.append(item[0])
    else:
        log.error("selectMode error: select from [define,min,max]")

    # 2.1、选取在selectPersonIds中的簇
    dataResList = []
    print("=====start select cluster=====")
    with tqdm(total=len(selectPersonIds)) as bar:
        for personId in selectPersonIds:
            dataResList.append(data[data["gt_person_id"] == personId])
            bar.update(1)
    print("=====end select cluster=====")

    log.info("cluster cluster finished and save result")
    # 2.2、存储筛选结果
    savePath = r"%s/%s_select_%s.pkl" % (dataDirPath, gtDataFileName, selectMode)
    print("簇筛选结果存储路径：%s" % savePath)
    write_pkl(savePath, pd.concat(dataResList))


if __name__ == '__main__':
    main_select_cluster()
